import logging
import re
import json
import asyncio
from typing import List, Dict, Optional
import chat

# --- 配置 ---
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)
logger = logging.getLogger(__name__)


# --- 辅助函数 ---

def split_text(text: str, max_length: int = 10000) -> List[str]:
    """
    将长文本按句号分割成不超过最大长度的块。
    """
    if len(text) <= max_length:
        return [text]

    chunks = []
    current_chunk_start = 0
    while current_chunk_start < len(text):
        end_pos = current_chunk_start + max_length
        if end_pos >= len(text):
            chunks.append(text[current_chunk_start:])
            break

        split_pos = text.rfind('。', current_chunk_start, end_pos)

        if split_pos == -1:
            split_pos = end_pos

        chunks.append(text[current_chunk_start:split_pos + 1])
        current_chunk_start = split_pos + 1

    logger.info(f"长文本被分割成 {len(chunks)} 个块进行处理。")
    return chunks


def build_extraction_prompt(text_chunk: str, instruction: Optional[str] = None) -> str:
    """
    构建一个强大的指令，要求模型一次性提取实体、类型、描述和关系。
    """
    instruction_section = f"**用户的额外要求 (最高优先级)**:\n{instruction}" if instruction else "请全面地识别所有实体、类型、描述和关系。"

    return f"""你是一个顶级的知识图谱构建专家。你的任务是从我提供的文本中，提取出所有实体、它们的通用类型、类型的描述，以及实体之间的关系。

请严格遵循以下规则：

1.  **遵守用户要求**：{instruction_section},且输出的格式也不能改变。
2.  **识别实体**: 识别文本中所有具体的实体名称，如“北京木材公司”、“张建国”、“《经营管理办法》”、“2023年3月1日”、“河北”等。
3.  **分配类型**: 为每个实体分配一个最合适的通用类型标签，如“公司”、“人员”、“文件”、”时间”、“地点”等。
4.  **生成类型描述**: 为每个**实体类型**（不是实体实例）生成一句简洁的通用描述。
5.  **推断关系 (Edges)**: 找出**具体实体**之间的关系三元组。关系应是简洁的动词或动词短语，避免使用“是”、“有”等无意义词汇。

6.  **输出格式**: 你的回复必须是一个严格的JSON字符串对象，且只包含以下两个键：`"nodes"` 和 `"edges"`。
    -   `"nodes"`: 一个对象数组。每个对象代表一个识别出的实体，必须包含三个键：
        -   `"id"`: (字符串) 实体在文本中的具体名称。
        -   `"type"`: (字符串) 该实体的通用类型。
        -   `"description"`: (字符串) 该实体**类型**的通用描述。
    -   `"edges"`: 一个对象数组。每个对象代表一个关系，必须包含三个键：
        -   `"source"`: (字符串) 源实体的具体名称。
        -   `"relation"`: (字符串) 它们之间的关系。
        -   `"target"`: (字符串) 目标实体的具体名称。

7.  **【至关重要的要求】**：
    -   你的回复必须只能包含JSON文本，不要有任何解释或代码块标记，输出要严格按照**输出格式**进行输出。
    -   你的回复必须直接以 `{{` 字符开始，并以 `}}` 字符结束。
    -   如果文本中没有实体或关系，返回包含空数组的JSON：`{{"nodes": [], "edges": []}}`。
    -   提取实体尤其注意要全面，尽可能地包含所有信息，实体自身和自身不能构建关系。

现在，请处理以下文本内容：

---
{text_chunk}
---
"""


def call_chat_api_non_stream(messages: List[Dict]) -> Optional[str]:
    """
    一个辅助函数，用于以非流式方式调用chat.py中的问答API。
    """
    try:
        response = chat.call_llm_api(messages, stream=False)
        if response and isinstance(response, dict):
            return response.get("choices", [{}])[0].get("message", {}).get("content", "")
        return None
    except Exception as e:
        logger.error(f"执行非流式API调用时发生错误: {e}", exc_info=True)
        return None


async def _process_chunk_async(text_chunk: str, instruction: Optional[str] = None) -> Dict:
    """
    异步处理单个文本块的提取任务，并传入instruction。
    """
    prompt = build_extraction_prompt(text_chunk, instruction)
    messages = [{"role": "user", "content": prompt}]

    model_response = await asyncio.to_thread(call_chat_api_non_stream, messages)

    if not model_response:
        return {"nodes": [], "edges": []}

    try:
        json_match = re.search(r'\{.*}', model_response, re.DOTALL)
        if not json_match: raise ValueError("未找到JSON对象结构")

        data = json.loads(json_match.group(0))
        if isinstance(data, dict) and "nodes" in data and "edges" in data:
            return data
        else:
            raise TypeError("解析出的JSON格式不符合要求")
    except (json.JSONDecodeError, ValueError, TypeError) as e:
        logger.error(f"解析模型返回的JSON时失败: {e}", extra={'raw_model_output': model_response})
        return {"nodes": [], "edges": []}


# --- 主处理函数 ---

async def extract_all_entities_and_relationships(text: str, instruction: Optional[str] = None) -> List[Dict]:
    """
    主处理函数：处理长文本，并发调用模型，合并并精确重组为新的扁平化格式。
    """
    if not text:
        return []

    # 1. 将文本分块
    text_chunks = split_text(text)

    # 2. 并发处理所有块
    tasks = [_process_chunk_async(chunk, instruction) for chunk in text_chunks]
    chunk_results = await asyncio.gather(*tasks)

    # 3. 创建一个从实体名称到其类型和描述的“速查表”
    node_map = {}
    for result in chunk_results:
        for node in result.get("nodes", []):
            node_id = node.get("id")
            if node_id and node_id not in node_map:
                node_map[node_id] = {
                    "type": node.get("type", "未知类型"),
                    "description": node.get("description", "无描述")
                }

    # 4. 遍历所有关系，并使用“速查表”来构建最终的扁平化列表
    final_data = []
    edge_tuples = set()  # 用于确保最终关系不重复

    for result in chunk_results:
        for edge in result.get("edges", []):
            source_id = edge.get("source")
            target_id = edge.get("target")
            relation = edge.get("relation")

            # 检查关系是否完整且未被添加过
            edge_tuple = (source_id, relation, target_id)
            if not all(edge_tuple) or edge_tuple in edge_tuples:
                continue

            # 从速查表中查找源和目标实体的信息
            source_node_info = node_map.get(source_id)
            target_node_info = node_map.get(target_id)

            # 只有当源和目标实体都存在于速查表中时，才认为这是一条合法的关系
            if source_node_info and target_node_info:
                # 构建您要求的扁平化字典结构
                final_data.append({
                    "rel": relation,
                    "source": source_id,
                    "sou_type": source_node_info["type"],
                    "sou_des": source_node_info["description"],
                    "target": target_id,
                    "tar_type": target_node_info["type"],
                    "tar_des": target_node_info["description"]
                })

                # 将该关系标记为已处理
                edge_tuples.add(edge_tuple)

    return final_data


#  功能2
def parse_instruction_for_more_type(instruction: str) -> List[Dict]:
    """
    解析 "类型-关系-类型(((描述)))" 格式的字符串，返回结构化的模板列表。
    """
    parsed_templates = []
    # 按逗号分割成独立的模板条目
    templates = [t.strip() for t in instruction.split(',') if t.strip()]

    # 正则表达式，用于匹配 "类型-关系-类型" 部分和可选的描述
    pattern = re.compile(r'(.+?)-(.+?)-(.+?)(?:\({3}(.*?)\){3})?$')

    for tpl_str in templates:
        match = pattern.match(tpl_str)
        if match:
            source_type, relation, target_type, description = match.groups()
            parsed_templates.append({
                "full": tpl_str,
                "source": source_type.strip(),
                "target": target_type.strip(),
                "description": description.strip() if description else None
            })
    return parsed_templates


def build_more_type_prompt(text_chunk: str, numbered_templates: List[str]) -> str:
    """
    构建一个指令，要求模型为每个带编号的模板寻找多个关系。
    """
    templates_str = "\n".join(numbered_templates)

    return f"""你是一个高级的语义关系分析专家。你的任务是根据我提供的“带编号的实体类型对列表”，在文本中为每一对实体类型找出最多5个、最少1个最相关的关系。

**带编号的实体类型对列表**:
---
{templates_str}
---

请严格遵循以下规则：

1.  **为每个编号提取**：你必须为列表中的**每一个**编号进行独立的分析。
2.  **利用描述**：如果模板中括号内有描述，请优先根据描述的上下文来寻找关系。
3.  **推断关系**：分析文本，找到连接源实体类型和目标实体类型的通用动词或上下文，形成关系。
4.  **数量要求**：为每个编号找出最多**5个**最贴切的关系，不得少于一个，提取要全面。
5.  **输出格式**：你的回复必须是纯文本，每行一个结果，并严格遵循 `编号: ["关系1", "关系2", ...]` 的JSON数组格式。

6.  **【至关重要的要求】**：
    -   **绝对不要添加任何解释、说明或代码块标记**。
    -   你的回复必须直接以第一个结果开始。
    -   如果某个编号在文本中找不到任何关系，请不要为它生成对应的行。
    -   示例输出:
1: ["召开", "举办"]
2: ["属于"]

现在，请处理以下文本内容：
---
{text_chunk}
---
"""


async def _process_chunk_for_more_type_async(text_chunk: str, numbered_templates: List[str]) -> Dict[str, List[str]]:
    """
    异步处理单个文本块的多关系提取任务。
    """
    prompt = build_more_type_prompt(text_chunk, numbered_templates)
    messages = [{"role": "user", "content": prompt}]

    model_response = await asyncio.to_thread(call_chat_api_non_stream, messages)

    chunk_relations = {}
    if model_response:
        # 正则表达式，用于匹配 "编号: [...]" 格式
        pattern = re.compile(r'^(\d+):\s*(\[.*?])', re.MULTILINE)
        matches = pattern.finditer(model_response)
        for match in matches:
            num, json_str = match.groups()
            try:
                relations = json.loads(json_str)
                if isinstance(relations, list):
                    chunk_relations[num] = relations
            except json.JSONDecodeError:
                logger.warning(f"在块处理中，未能解析编号 {num} 的JSON数组: {json_str}")
    return chunk_relations


async def extract_multiple_relationships(text: str, instruction: str) -> List[Dict]:
    """
    【新增】主处理函数：处理长文本，并发调用模型，合并并格式化多关系结果。
    """
    if not text or not instruction:
        return []

    # 1. 解析 instruction，为每个模板编号
    parsed_templates = parse_instruction_for_more_type(instruction)
    if not parsed_templates:
        logger.warning("从instruction中未能解析出任何有效的关系模板。")
        return []

    numbered_templates_for_prompt = []
    number_to_template_map = {}
    for i, tpl in enumerate(parsed_templates, 1):
        # prompt中只包含核心信息
        prompt_entry = f"[{i}] 源类型: {tpl['source']}, 目标类型: {tpl['target']}"
        if tpl['description']:
            prompt_entry += f" (描述: {tpl['description']})"
        numbered_templates_for_prompt.append(prompt_entry)
        number_to_template_map[str(i)] = tpl

    # 2. 将文本分块
    text_chunks = split_text(text)

    # 3. 并发处理所有块
    tasks = [_process_chunk_for_more_type_async(chunk, numbered_templates_for_prompt) for chunk in text_chunks]
    logger.info(f"准备并发执行 {len(tasks)} 个批次的多关系提取任务...")
    chunk_results = await asyncio.gather(*tasks)
    logger.info("所有并发任务已完成。")

    # 4. 合并并去重所有结果
    aggregated_relations = {}
    for result_dict in chunk_results:
        for num, relations in result_dict.items():
            if num not in aggregated_relations:
                aggregated_relations[num] = set()
            # 使用集合来自动去重
            for rel in relations:
                aggregated_relations[num].add(rel)

    # 5. 格式化为最终输出
    final_data = []
    # 遍历原始的编号映射表，以确保输出顺序和完整性
    for num, tpl in number_to_template_map.items():
        relations = list(aggregated_relations.get(num, []))[:5]  # 最多取5个
        if not relations:  # 如果一个关系都找不到，则跳过
            continue

        # 构建最终的输出对象
        output_obj = {
            "source": tpl["source"],
            "target": tpl["target"]
        }
        # 将关系动态地添加到对象中
        rel_keys = ["rel_one", "rel_two", "rel_three", "rel_four", "rel_five"]
        for i, rel in enumerate(relations):
            output_obj[rel_keys[i]] = rel

        final_data.append(output_obj)

    return final_data